#include "common.h"
#include "memory.h"
#include <string.h>
#include <elf.h>

#define ELF_OFFSET_IN_DISK 0

#ifdef HAS_DEVICE
void ide_read(uint8_t*, uint32_t, uint32_t);
#else
void ramdisk_read(uint8_t*, uint32_t, uint32_t);
#endif

#define STACK_SIZE (1 << 20)

void create_video_mapping();
uint32_t get_ucr3();

uint32_t loader() {
    Elf32_Ehdr* elf;
    Elf32_Phdr* ph = NULL;

    uint8_t buf[4096];

#ifdef HAS_DEVICE
    ide_read(buf, ELF_OFFSET_IN_DISK, 4096);
#else
    ramdisk_read(buf, ELF_OFFSET_IN_DISK, 4096);
#endif

    elf = (void*)buf;

    /* TODO: fix the magic number with the correct one */
    const uint32_t elf_magic = 0x464c457f;
    uint32_t* p_magic = (void*)buf;
    nemu_assert(*p_magic == elf_magic);

    uint32_t ph_num = elf->e_phnum;
    uint32_t ph_offset = elf->e_phoff;
    uint32_t ph_entry_size = elf->e_phentsize;

    /* Load each program segment */
    uint32_t i;
    for (i = 0; i < ph_num; i++) {
        /* Scan the program header table, load each segment into memory */
        uint8_t ph_buf[sizeof(Elf32_Phdr)];
#ifdef HAS_DEVICE
        ide_read(ph_buf, ph_offset + i * ph_entry_size, sizeof(Elf32_Phdr));
#else
        ramdisk_read(ph_buf, ph_offset + i * ph_entry_size, sizeof(Elf32_Phdr));
#endif
        
        ph = (Elf32_Phdr*)ph_buf;

        if (ph->p_type == PT_LOAD) {

            /* TODO: read the content of the segment from the ELF file
             * to the memory region [VirtAddr, VirtAddr + FileSiz)
             */
            uint32_t vaddr = ph->p_vaddr;
            uint32_t filesz = ph->p_filesz;
            uint32_t memsz = ph->p_memsz;

             /* TODO: zero the memory region
              * [VirtAddr + FileSiz, VirtAddr + MemSiz)
              */
            if (memsz > filesz) {
                memset((void*)(vaddr + filesz), 0, memsz - filesz);
            }

#ifdef IA32_PAGE
              /* Record the program break for future use. */
            extern uint32_t cur_brk, max_brk;
            uint32_t new_brk = ph->p_vaddr + ph->p_memsz - 1;
            if (cur_brk < new_brk) { max_brk = cur_brk = new_brk; }
#endif
        }
    }

    volatile uint32_t entry = elf->e_entry;

#ifdef IA32_PAGE
    mm_malloc(KOFFSET - STACK_SIZE, STACK_SIZE);

#ifdef HAS_DEVICE
    create_video_mapping();
#endif

    write_cr3(get_ucr3());
#endif

    return entry;
}
